obl=80; vary='S0'; vary_range=[1200 1850]; nyr=(vary_range(2)-vary_range(1))/10;
%obl=0; vary='S0'; vary_range=[1365 1850]; nyr=ceil((vary_range(2)-vary_range(1))/10);
obl=80; vary='Om'; vary_range=[5 0.2]; nyr=80; % only have 76 years...continue run
obl=0; vary='Om'; vary_range=[5 0.2]; nyr=80;

season='ANN';

casename=['OBL',sprintf('%d',obl),'_',vary];
head='vinth3_All_monthly_lonavg_';
filename_=[head,casename,'*_combine.nc'];
filename=dir(filename_);
filename=filename.name
crange=1e-8;
crange1=100;
showtop=50;showbot=950;
switch season
    case 'ANN' 
        months=[1:12];
    case 'Jan'
        months=[1];
    case 'DJF'
        months=[1 2 12];
    case 'MAM'
        months=[3 4 5];
end
plt_laplace=false;
plt_psi=true;
pltline=true;
diro='./decompose_streamfunction_transient/';
diroplt='./decompose_streamfunction_transient/';
%diroplt='~/Desktop/';

%% construct a varying variable series
yrs=[1:nyr];
if strcmp(vary,'S0')
    vary_ts=vary_range(1)+10.*yrs;
end
if strcmp(vary,'Om')
    vary_ts=vary_range(1).*(vary_range(2)/vary_range(1)).^((yrs-1)./80);
end
%% read
% coordinate
global nlat nlev lat lev clat S f R a mask U f0
lat=ncread(filename,'lat');
nlat=length(lat);
lev=ncread(filename,'lev');
ilev=ncread(filename,'ilev');
dp=ilev(2:end)-ilev(1:end-1);
nlev=length(lev);
time=ncread(filename,'time');
ntime=length(time);


% read monthly data and catagorize it into different months
VUs=squeeze(ncread(filename,'VU',[1,1,1],[Inf,Inf,Inf])); VUs=reshape(VUs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
VTs=squeeze(ncread(filename,'VT',[1,1,1],[Inf,Inf,Inf])); VTs=reshape(VTs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
Vs=squeeze(ncread(filename,'V',[1,1,1],[Inf,Inf,Inf])); Vs=reshape(Vs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
Us=squeeze(ncread(filename,'U',[1,1,1],[Inf,Inf,Inf])); Us=reshape(Us(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
Ts=squeeze(ncread(filename,'T',[1,1,1],[Inf,Inf,Inf])); Ts=reshape(Ts(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
RADs=squeeze(ncread(filename,'QRL',[1,1,1],[Inf,Inf,Inf]))./86400....
    +squeeze(ncread(filename,'QRS',[1,1,1],[Inf,Inf,Inf]))./86400; RADs=reshape(RADs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
CONVs=squeeze(ncread(filename,'DTCOND',[1,1,1],[Inf,Inf,Inf])); CONVs=reshape(CONVs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
PTTENDs=squeeze(ncread(filename,'PTTEND',[1,1,1],[Inf,Inf,Inf])); PTTENDs=reshape(PTTENDs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
UAPs=squeeze(ncread(filename,'UAP',[1,1,1],[Inf,Inf,Inf])); UAPs=reshape(UAPs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
OMEGAs=squeeze(ncread(filename,'OMEGA',[1,1,1],[Inf,Inf,Inf])); OMEGAs=reshape(OMEGAs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
OMEGAUs=squeeze(ncread(filename,'OMEGAU',[1,1,1],[Inf,Inf,Inf])); OMEGAUs=reshape(OMEGAUs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
OMEGATs=squeeze(ncread(filename,'OMEGAT',[1,1,1],[Inf,Inf,Inf])); OMEGATs=reshape(OMEGATs(:,:,1:12*nyr),[nlat,nlev,12,nyr]);
PSs=squeeze(ncread(filename,'PS',[1,1],[Inf,Inf])); PSs=reshape(PSs(:,1:12*nyr),[nlat,12,nyr]);

% yearly data for the interested months
VUs=squeeze(mean(VUs(:,:,months,:),3));
VTs=squeeze(mean(VTs(:,:,months,:),3));
Vs=squeeze(mean(Vs(:,:,months,:),3));
Us=squeeze(mean(Us(:,:,months,:),3));
Ts=squeeze(mean(Ts(:,:,months,:),3));
RADs=squeeze(mean(RADs(:,:,months,:),3));
CONVs=squeeze(mean(CONVs(:,:,months,:),3));
PTTENDs=squeeze(mean(PTTENDs(:,:,months,:),3));
UAPs=squeeze(mean(UAPs(:,:,months,:),3));
OMEGAs=squeeze(mean(OMEGAs(:,:,months,:),3));
OMEGAUs=squeeze(mean(OMEGAUs(:,:,months,:),3));
OMEGATs=squeeze(mean(OMEGATs(:,:,months,:),3));
PSs=squeeze(mean(PSs(:,months,:),2));

% cluster years
navg=3;
nskip=2; % 2 means every other cluster, 1 means every cluster etc.
samp=[5:navg*nskip:nyr];
VUs=movmean(VUs,navg,3); VUs=VUs(:,:,samp);
VTs=movmean(VTs,navg,3); VTs=VTs(:,:,samp);
Vs=movmean(Vs,navg,3); Vs=Vs(:,:,samp);
Us=movmean(Us,navg,3); Us=Us(:,:,samp);
Ts=movmean(Ts,navg,3); Ts=Ts(:,:,samp);
RADs=movmean(RADs,navg,3); RADs=RADs(:,:,samp);
CONVs=movmean(CONVs,navg,3); CONVs=CONVs(:,:,samp);
PTTENDs=movmean(PTTENDs,navg,3); PTTENDs=PTTENDs(:,:,samp);
UAPs=movmean(UAPs,navg,3); UAPs=UAPs(:,:,samp);
OMEGAs=movmean(OMEGAs,navg,3); OMEGAs=OMEGAs(:,:,samp);
OMEGAUs=movmean(OMEGAUs,navg,3); OMEGAUs=OMEGAUs(:,:,samp);
OMEGATs=movmean(OMEGATs,navg,3); OMEGATs=OMEGATs(:,:,samp);
PSs=movmean(PSs,navg,2); PSs=PSs(:,samp);
vary_ts=movmean(vary_ts,navg); vary_ts=vary_ts(samp);
nsamp=length(samp);

lev_=repmat(lev'.*100,[nlat,1,nsamp]);
PS_=repmat(reshape(PSs,[nlat,1,nsamp]),[1,nlev,1]);
masks=(lev_>PS_);
masks(1,:,:)=true;masks(end,:,:)=true;

VUs(1,:,:)=0.;VUs(end,:,:)=0.;
VTs(1,:,:)=0.;VTs(end,:,:)=0.;
Vs(1,:,:)=0.;Vs(end,:,:)=0.;


for isamp=1:nsamp
isamp,nsamp
%% get data for that sample
VU=squeeze(VUs(:,:,isamp));
VT=squeeze(VTs(:,:,isamp));
V=squeeze(Vs(:,:,isamp));
U=squeeze(Us(:,:,isamp));
T=squeeze(Ts(:,:,isamp));
RAD=squeeze(RADs(:,:,isamp));
CONV=squeeze(CONVs(:,:,isamp));
PTTEND=squeeze(PTTENDs(:,:,isamp));
UAP=squeeze(UAPs(:,:,isamp));
OMEGA=squeeze(OMEGAs(:,:,isamp));
OMEGAU=squeeze(OMEGAUs(:,:,isamp));
OMEGAT=squeeze(OMEGATs(:,:,isamp));
PS=squeeze(PSs(:,isamp));
vary_t=vary_ts(isamp);
mask=masks(:,:,isamp);

%% calculate omega equation terms
vt_eddy=VT-V.*T;
omegat_eddy=OMEGAT-OMEGA.*T;
uv_eddy=VU-V.*U;
omegau_eddy=OMEGAU-U.*OMEGA;

f0=2*2*pi/86400*vary_ts(isamp);
g=9.81;
a=6371000;
R=287;
kappa=2./7.;
f=repmat(reshape(f0*sind(lat),[nlat,1]),[1,nlev]);
clat=repmat(reshape(cosd(lat),[nlat,1]),[1,nlev]);
i_lat=[1:nlat];
p_lat=min(i_lat+1,nlat);
m_lat=max(i_lat-1,1);
dlat=repmat(reshape(lat(p_lat)-lat(m_lat),[nlat,1]),[1,nlev]).*(pi/180);
i_lev=[1:nlev];
p_lev=min(i_lev+1,nlev);
m_lev=max(i_lev-1,1);
z=log(1000./lev);
dz=repmat(reshape(z(p_lev)-z(m_lev),[1,nlev]),[nlat,1]);
dlev=repmat(reshape(lev(p_lev)-lev(m_lev),[1,nlev]),[nlat,1]).*100;
Th_T=(1000./lev').^kappa;

tmp=1./(clat.^2).*(uv_eddy(p_lat,:).*clat(p_lat,:).^2 - uv_eddy(m_lat,:).*clat(m_lat,:).^2)./dlat./a;
%tmp=1./(clat).*(uv_eddy(p_lat,:).*clat(p_lat,:) - uv_eddy(m_lat,:).*clat(m_lat,:))./dlat./a;
Term_uv=-f.*(tmp(:,p_lev)-tmp(:,m_lev))./dz;
Term_uv=Term_uv.*clat;

UV_y=1./clat.*V.*(U(p_lat,:).*clat(p_lat,:)-U(m_lat,:).*clat(m_lat,:))./dlat./a;
Term_UVadv=-f.*(UV_y(:,p_lev)-UV_y(:,m_lev))./dz;
Term_UVadv=Term_UVadv.*clat;

tmp=(omegau_eddy(:,p_lev) - omegau_eddy(:,m_lev))./dlev;
Term_uomega=-f.*(tmp(:,p_lev)-tmp(:,m_lev))./dz;
Term_uomega=Term_uomega.*clat.^2;
%Term_uomega(:)=0;

tmp=OMEGA.*(U(:,p_lev) - U(:,m_lev))./dlev;
Term_WUadv=-f.*(tmp(:,p_lev)-tmp(:,m_lev))./dz;
Term_WUadv=Term_WUadv.*clat;
%Term_WUadv(:)=0;

Term_CONV=R.*(CONV(p_lat,:)-CONV(m_lat,:))./dlat./a.*clat;
Term_RAD=R.*(RAD(p_lat,:)-RAD(m_lat,:))./dlat./a.*clat;
Term_Q=R.*(PTTEND(p_lat,:)-PTTEND(m_lat,:))./dlat./a.*clat;
%Term_Q=Term_CONV+Term_RAD;

dtime=3600;
FU=(UAP-U)./dtime;
Term_FU=f.*(FU(:,p_lev)-FU(:,m_lev))./dz;
Term_FU=Term_FU.*clat;

tmp=1./clat.*(vt_eddy(p_lat,:).*clat(p_lat,:) - vt_eddy(m_lat,:).*clat(m_lat,:))./dlat./a;
tmp(1,:)=0.;tmp(end,:)=0.;
Term_vt=-R.*(tmp(p_lat,:)-tmp(m_lat,:))./dlat./a;
Term_vt=Term_vt.*clat;

tmp=V.*(T(p_lat,:) - T(m_lat,:))./dlat./a;
Term_VTadv=-R.*(tmp(p_lat,:)-tmp(m_lat,:))./dlat./a;
Term_VTadv=Term_VTadv.*clat;

tmp=1./Th_T.*(omegat_eddy(:,p_lev).*Th_T(p_lev)-omegat_eddy(:,m_lev).*Th_T(m_lev))./dlev;
Term_omegat=-R.*(tmp(p_lat,:)-tmp(m_lat,:))./dlat./a;
Term_omegat=Term_omegat.*clat;

tmp=1./Th_T.*OMEGA.*(T(:,p_lev).*Th_T(p_lev)-T(:,m_lev).*Th_T(m_lev))./dlev;
Term_WTadv=-R.*(tmp(p_lat,:)-tmp(m_lat,:))./dlat./a;

Term_ADV=Term_VTadv+Term_WUadv+Term_UVadv;

Th=T.*(Th_T);
S=-T.*(log(Th(:,p_lev))-log(Th(:,m_lev)))./dlev;
Term_stm=f.^2.*(V(:,p_lev)-V(:,m_lev))./dz + R.*(S(p_lat,:).*OMEGA(p_lat,:)-S(m_lat,:).*OMEGA(m_lat,:))./dlat./a;
Term_stm=-Term_stm.*clat;

Term_uv(mask)=0.;
Term_UVadv(mask)=0.;
Term_uomega(mask)=0.;
Term_WUadv(mask)=0.;
Term_Q(mask)=0.;
Term_CONV(mask)=0.;
Term_RAD(mask)=0.;
Term_FU(mask)=0.;
Term_vt(mask)=0.;
Term_VTadv(mask)=0.;
Term_omegat(mask)=0.;
Term_WTadv(mask)=0.;
Term_ADV(mask)=0.;
Term_stm(mask)=0.;

%% plot omega equation terms
pltall=true;
if plt_laplace
if pltall
    totalplt=10;
else
    totalplt=5;
end

figure
polarmap;
iplt=1;
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_uv',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange+crange/20 crange-crange/20])
set(gca,'yscale','log','ydir','reverse')
title('uv')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
ylabel('Pressure (mb)')
xlabel('latitude')
iplt=iplt+1;

text(-150,35,[vary,'=',sprintf('%2.1f',vary_t)],'FontSize',24)

subplot(1,totalplt,iplt)
contourf(lat,lev,Term_vt',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('vt')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat,lev,Term_uomega',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('u\omega')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

if pltall
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_omegat',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('\omega t')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_CONV',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('convective heating')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_Q'-Term_CONV',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('other heating')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_FU',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('friction')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_ADV',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('Advection')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat,lev,Term_uv'+Term_uomega'+Term_Q'+Term_vt'+Term_omegat'+Term_ADV'+Term_FU',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('Sum')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
else
subplot(1,totalplt,iplt)
contourf(lat,lev,Term_uv'+Term_uomega'+Term_vt',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('Sum')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;
end

subplot(1,totalplt,iplt)
contourf(lat,lev,Term_stm',[-crange*10,[-crange-crange/20:crange/10:crange+crange/20],crange*10],'LineColor','none')
caxis([-crange-crange/20 crange+crange/20])
set(gca,'yscale','log','ydir','reverse')
title('\nabla^2\Psi_m')
xticks([-90:30:90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;



hp4=[0.55    0.08    0.33    0.13];
colorbar('Position', [hp4(1)+hp4(3)+0.04  hp4(2)+ hp4(4)-0.02 0.007  hp4(4)*5],'Ticks',[-crange,-crange/2,0,crange/2,crange])

if pltall
    set(gcf,'Position',[237         633        3000         315])
    saveas(gcf,[diroplt,'/laplace_psi_',casename,sprintf('%d',isamp),'_',season,'_all.png'])
else
    set(gcf,'Position',[237         633        1800         315])
    saveas(gcf,[diroplt,'/laplace_psi_',casename,sprintf('%d',isamp),'_',season,'.png'])
end
end

%% solve poisson
if ~isfile([diro,'/psi_driver_',casename,'_',season,'_mask.mat']) 
 smt_lat=1;
 smt_lev=1;
 if isamp==1
     %load([diro,'psi_driver_',casename,'_',season,'_mask.mat'])
     [psi_eddyfric0,lat1,lev1]=poisson_stm(smooth2a(-Term_uv-Term_vt-Term_uomega-Term_omegat-Term_FU,smt_lat,smt_lev));
     [psi_eddyheat0,lat1,lev1]=poisson_stm(smooth2a(-Term_vt-Term_omegat,smt_lat,smt_lev));
     [psi_thermal0,lat1,lev1]=poisson_stm(smooth2a(-Term_Q,smt_lat,smt_lev));
     [psi_ADV0,lat1,lev1]=poisson_stm(smooth2a(-Term_ADV,smt_lat,smt_lev));
     [psi_eddyvert0,lat1,lev1]=poisson_stm(smooth2a(-Term_uomega-Term_omegat,smt_lat,smt_lev));
     window_tropo=exp(-lat.^2./200);
     Term_eddyfrictrop=(Term_uv+Term_vt+Term_uomega+Term_omegat+Term_FU).*window_tropo;    
     [psi_eddyfrictrop0,lat1,lev1]=poisson_stm(smooth2a(-Term_eddyfrictrop,smt_lat,smt_lev));
     [psi_omegat0,lat1,lev1]=poisson_stm(smooth2a(-Term_omegat,smt_lat,smt_lev));        
 else
     [psi_eddyfric0,lat1,lev1]=poisson_stm(smooth2a(-Term_uv-Term_vt-Term_uomega-Term_omegat-Term_FU,smt_lat,smt_lev),psi_eddyfric0);
     [psi_eddyheat0,lat1,lev1]=poisson_stm(smooth2a(-Term_vt-Term_omegat,smt_lat,smt_lev),psi_eddyheat0);
     [psi_thermal0,lat1,lev1]=poisson_stm(smooth2a(-Term_Q,smt_lat,smt_lev),psi_thermal0);
     [psi_ADV0,lat1,lev1]=poisson_stm(smooth2a(-Term_ADV,smt_lat,smt_lev),psi_ADV0);
     [psi_eddyvert0,lat1,lev1]=poisson_stm(smooth2a(-Term_uomega-Term_omegat,smt_lat,smt_lev),psi_eddyvert0);
     window_trop=exp(-lat.^2./200);
     Term_eddyfrictrop=(Term_uv+Term_vt+Term_uomega+Term_omegat+Term_FU).*window_trop;
     [psi_eddyfrictrop0,lat1,lev1]=poisson_stm(smooth2a(-Term_eddyfrictrop,smt_lat,smt_lev),psi_eddyfrictrop0);
     [psi_omegat0,lat1,lev1]=poisson_stm(smooth2a(-Term_omegat,smt_lat,smt_lev),psi_omegat0);   
  end
 
 nlev1=length(lev1);
 
%% meridional streamfunction
ilev=ncread(filename,'ilev');
ilev1=ilev(2:end);
dp=100.*(ilev(2:end)-ilev(1:end-1));
psi0=2*pi*a/g*cosd(lat).*cumsum(V.*dp',2);

%% units
unit=1/g*2*pi*a/1e9;
 psi_eddyfric=psi_eddyfric0(:,1:nlev1).*unit;
 psi_eddyheat=psi_eddyheat0(:,1:nlev1).*unit;
 psi_thermal=psi_thermal0(:,1:nlev1).*unit;
 psi_ADV=psi_ADV0(:,1:nlev1).*unit;
 psi_eddyvert=psi_eddyvert0(:,1:nlev1).*unit;
 psi_eddyfrictrop=psi_eddyfrictrop0(:,1:nlev1).*unit;
 psi_omegat=psi_omegat0(:,1:nlev1).*unit;
 
 psi1=psi0/1e9;
 [zz,yy]=meshgrid(ilev1,lat);
 [zz1,yy1]=meshgrid(lev1,lat1);
 psi=interp2(zz,yy,psi1,zz1,yy1);
 
 psi_eddyfrics(:,:,isamp)=psi_eddyfric;
 psi_eddyheats(:,:,isamp)=psi_eddyheat;
 psi_thermals(:,:,isamp)=psi_thermal;
 psi_ADVs(:,:,isamp)=psi_ADV;
 psis(:,:,isamp)=psi;
 psi_eddyverts(:,:,isamp)=psi_eddyvert;
 psi_eddyfrictrops(:,:,isamp)=psi_eddyfrictrop;
 psi_omegats(:,:,isamp)=psi_omegat;
 


 if isamp==nsamp
    psi_eddyfrics=psi_eddyfrics(:,:,1:nsamp);
    psi_eddyheats=psi_eddyheats(:,:,1:nsamp);
    psi_thermals=psi_thermals(:,:,1:nsamp);
    psi_ADVs=psi_ADVs(:,:,1:nsamp);
    psi_eddyverts=psi_eddyverts(:,:,1:nsamp);
    psi_eddyfrictrops=psi_eddyfrictrops(:,:,1:nsamp);
    psi_omegats=psi_omegats(:,:,1:nsamp);
    psis=psis(:,:,1:nsamp);
    save([diro,'psi_driver_',casename,'_',season,'_mask.mat'],'psis','psi_eddyfrics','psi_eddyheats','psi_thermals','psi_ADVs','psi_eddyfrics','psi_eddyfrictrops','psi_eddyverts','psi_omegats','lev1','lat1','vary_ts');
 end
else
 if isamp==1
  load([diro,'psi_driver_',casename,'_',season,'_mask.mat'])
 end
 nlev1=length(lev1);
 psi=psis(:,:,isamp);
 psi_thermal=psi_thermals(:,:,isamp);
 psi_eddyfric=psi_eddyfrics(:,:,isamp);
 psi_eddyheat=psi_eddyheats(:,:,isamp);
 psi_ADV=psi_ADVs(:,:,isamp);
 psi_eddyvert=psi_eddyverts(:,:,isamp);
 psi_eddyfrictrop=psi_eddyfrictrops(:,:,isamp);
 psi_omegat=psi_omegats(:,:,isamp);
 
 unit=1/g*2*pi*a/1e9;
 psi_thermal0=psi_thermal./unit;%psi_thermal0=[psi_thermal0,psi_thermal0(:,end)]; % for extra boundary layer
 psi_eddyfric0=psi_eddyfric./unit;%psi_eddyfric0=[psi_eddyfric0,psi_eddyfric0(:,end)];
 psi_eddyheat0=psi_eddyheat./unit;%psi_eddyheat0=[psi_eddyheat0,psi_eddyheat0(:,end)];
 psi_ADV0=psi_eddyheat./unit;%psi_ADV0=[psi_ADV0,psi_ADV0(:,end)];
 psi_eddyvert0=psi_eddyvert./unit;%psi_eddyvert0=[psi_eddyvert0,psi_eddyvert0(:,end)];
 psi_eddyfrictrop0=psi_eddyfrictrop./unit;%psi_eddyfrictrop0=[psi_eddyfrictrop0,psi_eddyfrictrop0(:,end)];
 psi_omegat0=psi_omegat./unit;%psi_omegat0=[psi_omegat0,psi_omegat0(:,end)];
end
%% plot 
if plt_psi
totalplt=5;
iplt=1;

figure
polarmap
% subplot(1,totalplt,iplt)
% contourf(lat1,lev1,psi_uv',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
% caxis([-crange1-crange1/20 crange1-crange1/20])
% set(gca,'yscale','log','ydir','reverse')
% title('\psi_{uv}')
% xticks([-90:30:90])
% xlim([-90 90])
% yticks([4,10,20,50,100,200,500,950])
% ylim([showtop,showbot])
% set(gca,'FontSize',20)
% xlabel('latitude')
% iplt=iplt+1;
% 
% subplot(1,totalplt,iplt)
% contourf(lat1,lev1,psi_vt',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
% caxis([-crange1-crange1/20 crange1-crange1/20])
% set(gca,'yscale','log','ydir','reverse')
% title('\psi_{vt}')
% xticks([-90:30:90])
% xlim([-90 90])
% yticks([4,10,20,50,100,200,500,950])
% ylim([showtop,showbot])
% set(gca,'FontSize',20)
% xlabel('latitude')
% iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_eddyfric'-psi_eddyheat',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
%title('uv+vt+u\omega+friction')
%title('adiabatic')
title('momentum')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

text(-220,38,[vary,'=',sprintf('%2.1f',vary_t)],'FontSize',24)

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_thermal'+psi_eddyheat',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
%title('thermal+t\omega')
%title('diabatic')
title('thermal')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_ADV',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('advection')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_eddyfric'+psi_thermal'+psi_ADV',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('Sum')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('\Psi')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

hp4=[0.55    0.08    0.33    0.13];
colorbar('Position', [hp4(1)+hp4(3)+0.04  hp4(2)+ hp4(4)-0.02 0.009  hp4(4)*5.5],'Ticks',[-crange1,-crange1/2,0,crange1/2,crange1])


set(gcf,'Position',[237         633        1800         315])
saveas(gcf,[diroplt,'/psi_decomp_momentum_',casename,sprintf('%d',isamp),'_',season,'.png'])
end

%%
if plt_psi
totalplt=5;
iplt=1;

figure
polarmap
% subplot(1,totalplt,iplt)
% contourf(lat1,lev1,psi_uv',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
% caxis([-crange1-crange1/20 crange1-crange1/20])
% set(gca,'yscale','log','ydir','reverse')
% title('\psi_{uv}')
% xticks([-90:30:90])
% xlim([-90 90])
% yticks([4,10,20,50,100,200,500,950])
% ylim([showtop,showbot])
% set(gca,'FontSize',20)
% xlabel('latitude')
% iplt=iplt+1;
% 
% subplot(1,totalplt,iplt)
% contourf(lat1,lev1,psi_vt',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
% caxis([-crange1-crange1/20 crange1-crange1/20])
% set(gca,'yscale','log','ydir','reverse')
% title('\psi_{vt}')
% xticks([-90:30:90])
% xlim([-90 90])
% yticks([4,10,20,50,100,200,500,950])
% ylim([showtop,showbot])
% set(gca,'FontSize',20)
% xlabel('latitude')
% iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_eddyfric',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
%title('uv+vt+u\omega+friction')
title('adiabatic')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

text(-220,38,[vary,'=',sprintf('%2.1f',vary_t)],'FontSize',24)

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_thermal',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
%title('thermal+t\omega')
title('diabatic')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_ADV',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('advection')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi_eddyfric'+psi_thermal'+psi_ADV',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('Sum')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

subplot(1,totalplt,iplt)
contourf(lat1,lev1,psi',[-crange1*10,[-crange1-crange1/20:crange1/10:crange1+crange1/20],crange1*10],'LineColor','none')
caxis([-crange1-crange1/20 crange1-crange1/20])
set(gca,'yscale','log','ydir','reverse')
title('\Psi')
xticks([-90:30:90])
xlim([-90 90])
yticks([4,10,20,50,100,200,500,950])
ylim([showtop,showbot])
set(gca,'FontSize',20)
xlabel('latitude')
iplt=iplt+1;

hp4=[0.55    0.08    0.33    0.13];
colorbar('Position', [hp4(1)+hp4(3)+0.04,  hp4(2)+ hp4(4)-0.02, 0.009  ,hp4(4)*5.5],'Ticks',[-crange1,-crange1/2,0,crange1/2,crange1])


set(gcf,'Position',[237         633        1800         315])
saveas(gcf,[diroplt,'/psi_decomp_adiabatic_',casename,sprintf('%d',isamp),'_',season,'.png'])
end

end

%% line plot
latrange=[-90 90];
levrange=[1000 1];
ilev1r=(lev1>levrange(2) & lev1<levrange(1));
lev1r=lev1(ilev1r);
ilat1r=(lat1>latrange(1) & lat1<latrange(2));
lat1r=lat1(ilat1r);
nlev1r=length(lev1r);nlat1r=length(lat1r);

max_eddyfric=nanmax(reshape(psi_eddyfrics(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_eddyheat=nanmax(reshape(psi_eddyheats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_thermal=nanmax(reshape(psi_thermals(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_ADV=nanmax(reshape(psi_ADVs(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_eddyvert=nanmax(reshape(psi_eddyverts(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_eddyfrictrop=nanmax(reshape(psi_eddyfrictrops(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);
max_omegat=nanmax(reshape(psi_omegats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),[],1);

std_eddyfric=nanstd(reshape(psi_eddyfrics(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_eddyheat=nanstd(reshape(psi_eddyheats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_thermal=nanstd(reshape(psi_thermals(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_ADV=nanstd(reshape(psi_ADVs(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_eddyvert=nanstd(reshape(psi_eddyverts(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_eddyfrictrop=nanstd(reshape(psi_eddyfrictrops(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);
std_omegat=nanstd(reshape(psi_omegats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),0,1);

dp=ilev
weight=repmat(dp(ilev1r),[nlat1r,1]);
base1=nanmean(reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]),1);
base=nanmean(reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1);
proj_eddyfric=nanmean(reshape(psi_eddyfrics(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_eddyheat=nanmean(reshape(psi_eddyheats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_thermal=nanmean(reshape(psi_thermals(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_ADV=nanmean(reshape(psi_ADVs(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_eddyvert=nanmean(reshape(psi_eddyverts(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_eddyfrictrop=nanmean(reshape(psi_eddyfrictrops(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;
proj_omegat=nanmean(reshape(psi_omegats(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*reshape(psis(ilat1r,ilev1r,:),[nlev1r*nlat1r,nsamp]).*weight,1)./base;

if pltline
figure
hold on
 plot(vary_ts,abs(smooth((proj_eddyfric-proj_eddyheat).*sqrt(base1))).^0.5.*sign(smooth((proj_eddyfric-proj_eddyheat).*sqrt(base1))),'b','LineWidth',2)
% plot(vary_ts,abs(smooth(proj_eddyheat.*sqrt(base1))).^0.5.*sign(smooth(proj_eddyheat.*sqrt(base1))),'m--','LineWidth',2)
 plot(vary_ts,abs(smooth((proj_thermal+proj_eddyheat).*sqrt(base1),7)).^0.5.*sign(smooth((proj_thermal+proj_eddyheat).*sqrt(base1),7)),'r','LineWidth',2)
 plot(vary_ts,abs(smooth(proj_ADV.*sqrt(base1),7)).^0.5.*sign(smooth(proj_ADV.*sqrt(base1),7)),'k')
%plot(vary_ts,smooth((proj_eddyfric)),'b','LineWidth',2)
%plot(vary_ts,smooth(proj_eddyheat),'m--','LineWidth',2)
%plot(vary_ts,smooth((proj_thermal)),'r','LineWidth',2)
%plot(vary_ts,smooth(proj_ADV),'k')
%legend({'uv+vt+u\omega+fric','t\omega','thermal+t\omega','advection'},'Location','northwest')
legend({'momentum','thermal','advection'},'Location','northwest')
%legend({'adiabatic','diabatic','advection'},'Location','northwest')
xlim([min(vary_ts),max(vary_ts)])
ylabel('$$(\Psi_{\mathrm{proj}})^{1/2}$$','Interpreter','latex')
if strcmp(vary,'S0')
xlabel('Insolation (W/m^2)')
set(gca,'FontSize',20)
end
if strcmp(vary,'Om')
set(gca,'xscale','log')
xticks([0.25 0.5 1 2 4])
xlabel('Rotation rate (Earth rotation)')
set(gca,'FontSize',20)
end
title(['Obliquity=',sprintf('%d',obl),', ',season])
saveas(gcf,[diroplt,'/projhadley_momentum_',casename,'_',season,'.png'])

figure
hold on
 plot(vary_ts,abs(smooth((proj_eddyfric).*sqrt(base1))).^0.5.*sign(smooth((proj_eddyfric-proj_eddyheat).*sqrt(base1))),'b','LineWidth',2)
% plot(vary_ts,abs(smooth(proj_eddyheat.*sqrt(base1))).^0.5.*sign(smooth(proj_eddyheat.*sqrt(base1))),'m--','LineWidth',2)
 plot(vary_ts,abs(smooth((proj_thermal).*sqrt(base1),7)).^0.5.*sign(smooth((proj_thermal+proj_eddyheat).*sqrt(base1),7)),'r','LineWidth',2)
 plot(vary_ts,abs(smooth(proj_ADV.*sqrt(base1),7)).^0.5.*sign(smooth(proj_ADV.*sqrt(base1),7)),'k')
legend({'adiabatic','diabatic','advection'},'Location','northwest')
xlim([min(vary_ts),max(vary_ts)])
ylabel('$$(\Psi_{\mathrm{proj}})^{1/2}$$','Interpreter','latex')
if strcmp(vary,'S0')
xlabel('Insolation (W/m^2)')
set(gca,'FontSize',20)
end
if strcmp(vary,'Om')
set(gca,'xscale','log')
xticks([0.25 0.5 1 2 4])
xlabel('Rotation rate (Earth rotation)')
set(gca,'FontSize',20)
end
title(['Obliquity=',sprintf('%d',obl),', ',season])
saveas(gcf,[diroplt,'/projhadley_adiabatic_',casename,'_',season,'.png'])

% figure
% hold on
% var_all=[smooth(proj_thermal.*sqrt(base));...
%     smooth((proj_eddyfric-proj_eddyfrictrop).*sqrt(base));...
%     smooth(proj_eddyfrictrop.*sqrt(base));...
%     smooth(proj_ADV.*sqrt(base))];
% area(vary_ts,var_all,)
% legend({'momentum','thermal','advection'})
% xlim([min(vary_ts),max(vary_ts)])
% if strcmp(vary,'S0')
% xlabel('Insolation (W)')
% ylabel('% variance')
% set(gca,'FontSize',20)
% end
% if strcmp(vary,'Om')
% xlabel('Rotation rate (Earth rotation)')
% ylabel('Percent variance')
% set(gca,'FontSize',20)
% end
% title(['Obliquity=',sprintf('%d',obl),', ',season])
% saveas(gcf,[diroplt,'/projhadley_',casename,'_',season,'.png'])
end

%% without extra boundary layer
% function [psi,lat1,lev1]=poisson_stm(T,psi_init)
% global nlat nlev lat lev clat S f R a U f0
% nbot=0;
% lev1=lev(1:end-nbot);
% pres=lev1.*100;
% P0=101325;
% %pres=cat(1,pres,P0);
% kappa=2/5;
% 
% lat1=lat;
% clat1=clat(:,1:end-nbot);%clat1=cat(2,clat1,clat1(:,end));
% seclat1=1./clat1;seclat1(1,:)=seclat1(2,:);seclat1(end,:)=seclat1(end-1,:);
% S1=S(:,1:end-nbot);%S1=cat(2,S1,S1(:,end));
% f1=f(:,1:end-nbot);%f1=cat(2,f1,f1(:,end));
% T=T(:,1:end-nbot);%T=cat(2,T,T(:,end));
% U1=U(:,1:end-nbot);%U1=cat(2,U1,U1(:,end));
% d2r=pi/180;
% 
% [Ucos_p,Ucos_y]=gradient(U1.*clat1,pres,sind(lat1));
% if ~exist('psi_init','var')
%     psi0=zeros(nlat,nlev-nbot);
% else
%     psi0=psi_init;
% end
% psi=psi0;
% dt0=1e8;
% dt=dt0;
% err=1;
% err_min=1;
% nfail=0;
% 
% 
% for it=1:1000000
% if err>1e-12
%     psi0=psi;
%     %dt=dt0/it^0.1;
%     [psi_p,psi_y]=gradient(psi,pres,lat1.*d2r);
%     [psi_pp,~]=gradient(psi_p,pres,lat1.*d2r);
%     [~,psi_yy]=gradient(psi_y.*seclat1,pres,lat1.*d2r);
%     tend=dt.*(f1.^2.*pres'.*psi_pp+R/a^2.*clat1.*psi_yy.*S1+T);
%     tend=tend-dt.*(-pres'.*f0./a^2.*Ucos_p.*psi_p+f1.*(1-kappa).*Ucos_p.*seclat1./a.*psi_y+f1.*pres'./a.*Ucos_y.*psi_pp);
%     tend(1,:)=0;tend(end,:)=0;tend(:,1)=0;tend(:,end)=0;
%     psi=psi0+tend;
%     
%     
%     if mod(it,1000)==0 | it<1000
%         err=std(tend(:)./dt);
%         if err<err_min
%             psi_save=psi;
%             dt=dt*1.00;
%             nfail=0;
%         else
%             dt=dt/1.2;
%             if floor(it/1000)-nfail<50
%                 psi=psi_save;
%             else
%                 psi=(psi_save+psi)./2;
%             end
%             nfail=nfail+1
%             if err/err_min>2 | nfail>100
%                 break;
%             end         
%         end
%         
%     if mod(it,1000)==0 
%         disp(it);disp(err)
%     end
%     err_min=min(err,err_min);
%     end
% end
% end
% 
% psi=psi_save;
% end

%% with extra boundary layer
% function [psi,lat1,lev1]=poisson_stm(T,psi_init)
% global nlat nlev lat lev clat S f R a U f0
% nbot=0;
% lev1=lev(1:end-nbot);
% pres=lev1.*100;
% P0=101325;
% kappa=2/5;
% pres=lev1.*100;pres=cat(1,pres,P0);
% lat1=lat;
% clat1=clat(:,1:end-nbot);clat1=cat(2,clat1,clat1(:,end));
% seclat1=1./clat1;seclat1(1,:)=seclat1(2,:);seclat1(end,:)=seclat1(end-1,:);
% S1=S(:,1:end-nbot);S1=cat(2,S1,S1(:,end));
% f1=f(:,1:end-nbot);f1=cat(2,f1,f1(:,end));
% T=T(:,1:end-nbot);T=cat(2,T,T(:,end));
% U1=U(:,1:end-nbot);U1=cat(2,U1,U1(:,end));
% d2r=pi/180;
% 
% [Ucos_p,Ucos_y]=gradient(U1.*clat1,pres,sind(lat1));
% if ~exist('psi_init','var')
%     psi0=zeros(nlat,nlev-nbot+1);
% else
%     psi0=psi_init;
% end
% psi=psi0;
% dt0=1e8;
% dt=dt0;
% err=1;
% err_min=1;
% nfail=0;
% 
% 
% for it=1:1000000
% if err>1e-12
%     psi0=psi;
%     %dt=dt0/it^0.1;
%     [psi_p,psi_y]=gradient(psi,pres,lat1.*d2r);
%     [psi_pp,~]=gradient(psi_p,pres,lat1.*d2r);
%     [~,psi_yy]=gradient(psi_y.*seclat1,pres,lat1.*d2r);
%     tend=dt.*(f1.^2.*pres'.*psi_pp+R/a^2.*clat1.*psi_yy.*S1+T);
%     tend=tend-dt.*(-pres'.*f0./a^2.*Ucos_p.*psi_p+f1.*(1-kappa).*Ucos_p.*seclat1./a.*psi_y+f1.*pres'./a.*Ucos_y.*psi_pp);
%     tend(1,:)=0;tend(end,:)=0;tend(:,1)=0;tend(:,end)=0;
%     psi=psi0+tend;
%     
%     if mod(it,1000)==0
%         err=std(tend(:)./dt);
%         if err<err_min
%             psi_save=psi;
%             dt=dt*1.002;
%             nfail=0;
%         else
%             dt=dt/1.2;
%             if floor(it,1000)-nfail<50
%                 psi=psi_save;
%             else
%                 psi=(psi_save+psi)./2;
%             end
%             nfail=nfail+1
%             if err/err_min>2 | nfail>100
%                 break;
%             end         
%         end
%         
%     disp(it);disp(err)
%     err_min=min(err,err_min);
%     end
% end
% end
% 
% psi=psi_save;
% end

%% use mask and pressure coordinate
%% Poisson solver
function [psi,lat1,lev1]=poisson_stm(T,psi_init)
global nlat nlev lat lev clat S f R a mask U f0
nbot=0;
lev1=lev(1:end-nbot);
P0=101325;
kappa=2/5;
pres=lev1.*100;%pres=cat(1,pres,P0);
lat1=lat;
clat1=clat(:,1:end-nbot);%clat1=cat(2,clat1,clat1(:,end));
seclat1=1./clat1;seclat1(1,:)=1.;seclat1(end,:)=1.;
S1=S(:,1:end-nbot);%S1=cat(2,S1,S1(:,end));
f1=f(:,1:end-nbot);%f1=cat(2,f1,f1(:,end));
T=T(:,1:end-nbot);%T=cat(2,T,T(:,end));
U1=U(:,1:end-nbot);%U1=cat(2,U1,U1(:,end));
ones1=mask(:,end);
ones1(:,:)=true;
mask1=cat(2,mask(:,2:end),ones1);
%mask1=mask;

d2r=pi/180;

[Ucos_p,Ucos_y]=gradient(U1.*clat1,pres,sind(lat1));

if ~exist('psi_init','var')
     %psi0=zeros(nlat,nlev-nbot+1);
     psi0=zeros(nlat,nlev-nbot);
else
     psi0=psi_init;
end
psi=psi0;

dt0=3e7;
dt=dt0;
err=1;
err_min=1;
err_previous=1;
nfail=0;


for it=1:3000000
if err>1e-12
    psi0=psi;
    %dt=dt0/(it)^0.1;
    [psi_p,psi_y]=gradient(psi,pres,lat1.*d2r);
    [psi_pp,~]=gradient(psi_p,pres,lat1.*d2r);
    [~,psi_yy]=gradient(psi_y.*S1.*seclat1,pres,lat1.*d2r);
    tend=dt.*(f1.^2.*pres'.*psi_pp+R/a^2.*clat1.*psi_yy+T);
    %tend=tend-dt.*(-pres'.*f0./a^2.*Ucos_p.*psi_p+f1.*(1-kappa).*Ucos_p.*seclat1./a.*psi_y+f1.*pres'./a.*Ucos_y.*psi_pp);
    tend(:,1)=0;
    
    %tend(1,:)=0;tend(end,:)=0;
    %tend(:,end)=0;
    
    tend(mask1)=0;
    
    psi=psi0+tend;
    
    if mod(it,1000)==0 | it<1000
        err=std(tend(:)./dt);
        if err<err_min
            psi_save=psi;
            dt=dt*1.00;
            nfail=0;
        elseif err<err_previous
            psi_save=psi;
            dt=dt*1.00;
            nfail=0;
        else
            dt=dt/1.2;
            if floor(it/1000)-nfail<50
                psi=psi_save;
            else
                psi=psi_save;
            end
            nfail=nfail+1
            if err/err_min>2 | nfail>100
                break;
            end         
        end
        
    if mod(it,1000)==0 
        disp(it);disp(err)
    end
    err_min=min(err,err_min);
    err_previous=err;
    end
end
end
psi=psi_save;
end